import wcocr
import os
import tempfile
import base64
from io import BytesIO
from docx import Document
from docx.shared import Pt
from docx.oxml.ns import qn
from colorama import init, Fore, Style

# 尝试导入PIL，如果没有安装则提示
try:
    from PIL import Image
    import numpy as np
    PIL_AVAILABLE = True
except ImportError:
    PIL_AVAILABLE = False
    print("提示：如需处理内存图片，请安装PIL: pip install Pillow")

def find_wechat_path():
    script_dir = os.path.dirname(os.path.abspath(__file__))
    common_paths = os.path.join(script_dir, 'path')
    if os.path.exists(common_paths):
        return common_paths
    else:
        print(f"The path folder does not exist at {common_paths}.")
        return None

def find_wechatocr_exe():
    script_dir = os.path.dirname(os.path.abspath(__file__))
    wechatocr_path = os.path.join(script_dir, 'path', 'WeChatOCR', 'WeChatOCR.exe')
    if os.path.isfile(wechatocr_path):
        return wechatocr_path
    else:
        print(f"The WeChatOCR.exe does not exist at {wechatocr_path}.")
        return None

def save_memory_image_to_temp(image_data, image_format='PNG'):
    """
    将内存中的图片数据保存为临时文件
    支持多种输入格式：PIL Image、numpy数组、base64字符串、字节数据
    """
    if not PIL_AVAILABLE:
        raise ImportError("处理内存图片需要安装PIL: pip install Pillow")
    
    # 创建临时文件
    temp_file = tempfile.NamedTemporaryFile(delete=False, suffix=f'.{image_format.lower()}')
    temp_path = temp_file.name
    temp_file.close()
    
    try:
        # 处理不同类型的输入
        if isinstance(image_data, str):
            # Base64字符串
            if image_data.startswith('data:image'):
                # 处理data URL格式
                image_data = image_data.split(',')[1]
            image_bytes = base64.b64decode(image_data)
            image = Image.open(BytesIO(image_bytes))
        elif isinstance(image_data, bytes):
            # 字节数据
            image = Image.open(BytesIO(image_data))
        elif isinstance(image_data, np.ndarray):
            # numpy数组
            image = Image.fromarray(image_data)
        elif hasattr(image_data, 'save'):
            # PIL Image对象
            image = image_data
        else:
            raise ValueError(f"不支持的图片数据类型: {type(image_data)}")
        
        # 保存为临时文件
        image.save(temp_path, format=image_format)
        return temp_path
        
    except Exception as e:
        # 清理临时文件
        if os.path.exists(temp_path):
            os.unlink(temp_path)
        raise e

def wechat_ocr_from_memory(image_data, image_format='PNG'):
    """
    从内存中的图片数据进行OCR识别
    """
    temp_path = None
    try:
        # 将内存图片保存为临时文件
        temp_path = save_memory_image_to_temp(image_data, image_format)
        
        # 使用现有的OCR函数处理
        result = wechat_ocr(temp_path)
        
        return result
        
    finally:
        # 清理临时文件
        if temp_path and os.path.exists(temp_path):
            os.unlink(temp_path)

def wechat_ocr(image_path):
    wechat_path = find_wechat_path()
    wechatocr_path = find_wechatocr_exe()
    if not wechat_path or not wechatocr_path:
        return []  # 返回空结果
    
    wcocr.init(wechatocr_path, wechat_path)
    result = wcocr.ocr(image_path)
    
    # 打印原始结果以便调试
    print("原始OCR结果:")
    print(result)
    
    ocr_results = []

    for temp in result['ocr_response']:
        text = temp['text']
        if isinstance(text, bytes):
            text = text.decode('utf-8', errors='ignore')
        
        # 提取完整信息包括坐标
        ocr_item = {
            'text': text,
            'coordinates': temp.get('coordinates', []),
            'left': temp.get('left', 0),
            'top': temp.get('top', 0),
            'right': temp.get('right', 0),
            'bottom': temp.get('bottom', 0)
        }
        ocr_results.append(ocr_item)
    
    return ocr_results

def save_to_docx(texts, output_path):
    doc = Document()

    for text in texts:
        # 添加段落并设置宋体字体
        paragraph = doc.add_paragraph()
        run = paragraph.add_run(text)
        run.font.name = '宋体'

        # 设置字体为宋体 (兼容中文设置)
        r = run._element
        r.rPr.rFonts.set(qn('w:eastAsia'), '宋体')

        # 设置字体大小为五号字体 (10.5 磅)
        run.font.size = Pt(10.5)
    
    doc.save(output_path)

def save_to_markdown(ocr_results, output_path, image_name):
    """保存OCR结果为Markdown格式，包含坐标信息"""
    with open(output_path, 'w', encoding='utf-8') as f:
        f.write(f"# OCR识别结果 - {image_name}\n\n")
        f.write("---\n\n")
        
        for i, item in enumerate(ocr_results, 1):
            f.write(f"## 文本块 {i}\n\n")
            f.write(f"**文本内容**: {item['text']}\n\n")
            
            # 添加坐标信息
            if item['coordinates']:
                f.write(f"**坐标信息**: {item['coordinates']}\n\n")
            
            if item['left'] or item['top'] or item['right'] or item['bottom']:
                f.write(f"**边界框**: 左:{item['left']}, 上:{item['top']}, 右:{item['right']}, 下:{item['bottom']}\n\n")
            
            f.write("---\n\n")
        
        f.write(f"*生成时间: {os.path.basename(output_path)}*\n")

def process_all_images():
    script_dir = os.path.dirname(os.path.abspath(__file__))
    src_folder = os.path.join(script_dir, 'src')
    md_folder = os.path.join(script_dir, 'markdown')

    if not os.path.exists(md_folder):
        os.makedirs(md_folder)

    # 支持的图像格式
    image_extensions = ('.jpg', '.jpeg', '.png', '.bmp', '.tif')

    # 遍历 src 文件夹及其所有子文件夹
    for root, dirs, files in os.walk(src_folder):
        for file in files:
            if file.lower().endswith(image_extensions):
                image_path = os.path.join(root, file)
                relative_path = os.path.relpath(root, src_folder)
                md_folder_path = os.path.join(md_folder, relative_path)

                # 确保 markdown 文件夹路径存在
                if not os.path.exists(md_folder_path):
                    os.makedirs(md_folder_path)

                # 处理图片文件
                print(Fore.GREEN + f"正在处理: {os.path.relpath(image_path, script_dir)}" + Style.RESET_ALL)
                ocr_results = wechat_ocr(image_path)
                image_name = os.path.splitext(file)[0]
                output_md = os.path.join(md_folder_path, f'{image_name}_OCR.md')
                save_to_markdown(ocr_results, output_md, image_name)
                # 显示相对路径
                relative_md_path = os.path.relpath(output_md, script_dir)
                print(f"OCR 结果已保存到： {relative_md_path}\n")

if __name__ == '__main__':
    init(autoreset=True)  # 初始化 colorama
    process_all_images()
    print(Fore.RED + "全部文件处理完成，请按 Enter 键退出……" + Style.RESET_ALL)
    input()
